package com.example.util;

import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.Callable;

public class MDCUtil {

    // 包装 Callable 任务，传递 MDC 上下文信息
    public static <T> Callable<T> wrapWithMdc(Callable<T> task) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return () -> {
            if (contextMap != null) {
                MDC.setContextMap(contextMap);
            }
            try {
                return task.call();
            } finally {
                MDC.clear();  // 确保上下文清理，避免污染其他任务
            }
        };
    }

    // 包装 Runnable 任务，传递 MDC 上下文信息
    public static Runnable wrapWithMdc(Runnable task) {
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        return () -> {
            if (contextMap != null) {
                MDC.setContextMap(contextMap);
            }
            try {
                task.run();
            } finally {
                MDC.clear();  // 清理 MDC 上下文
            }
        };
    }
}
